#!/usr/bin/env python
# LeCrunch
# Copyright (C) 2010 Anthony LaTorre 
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import time
import socket
import struct

import scope2
import rootfile

def main(filename, nevents, cfg_file=None):
    scope = scope2.ScopeReader()
    scope.connect()
    print 'Initializing scope...'
    sys.stdout.flush()

    # turn off the display
    scope.send('display off')
    scope.check_last_command()

    # clear the output queue
    scope.clear()

    if cfg_file is None:
        pass
    else:
        f = open(cfg_file)

        for line in f:
            msg = line.split('#')[0].strip()
            if msg != '':
                print 'Sending %s...' % repr(msg),
                scope.send(msg)
                if '?' in msg:
                    print
                    print '>> %s' % repr(scope.recv())
                scope.check_last_command()
                print 'Success!'

        f.close()

    channels = scope.getchannels()

    # prime the output queue
    for i in range(10):
        scope.trigger()
        for cnl in channels:
            scope.send('c%i:wf?' % cnl)

    # open up the output file
    output = rootfile.RootFile(filename, 'w')

    # start a timer
    time0 = time.time()

    try:
        try:
            for i in range(nevents):
                print '\rSaving event: %i' % (i+1),
                sys.stdout.flush()

                try:
                    scope.trigger()
                    if 1 in channels:
                        output.setc1(*scope.getwaveform(1))
                    if 2 in channels:
                        output.setc2(*scope.getwaveform(2))
                    if 3 in channels:
                        output.setc3(*scope.getwaveform(3))
                    if 4 in channels:
                        output.setc4(*scope.getwaveform(4))
                    output.fill()
                except (socket.error, struct.error):
                    print "\nFailed to receive waveform %i" % (i+1)
                    # clear the output queue
                    scope.clear()
                    # reprime the queue
                    for i in range(10):
                        scope.trigger()
                        for cnl in channels:
                            scope.send('c%i:wf?' % cnl)
                    continue
        except KeyboardInterrupt:
            print '\nTerminating...',
    finally:
        print

        elapsed = time.time() - time0

        print
        print 'Completed %i events in %.3f seconds.' % (i+1, elapsed)
        print 'Averaged %.5f seconds per acquisition.' % (elapsed/nevents)
        print "Wrote to file '%s'." % filename

        output.close()
        scope.send('display on')
        time.sleep(2.0)
        scope.close()

if __name__ == '__main__':
    import optparse

    usage = "usage: %prog <filename/prefix> [-n] [-r]"
    parser = optparse.OptionParser(usage, version="%prog 0.1.0")
    parser.add_option("-n", type="int", dest="nevents",
                      help="number of events to store per run", default=1000)
    parser.add_option("-r", type="int", dest="nruns",
                      help="number of runs", default=1)
    parser.add_option("-c", type="string", dest="cfg_file",
                      help="configuration file", default=None)
    (options, args) = parser.parse_args()

    if len(args) < 1:
        parser.print_help()
        sys.exit(0)
    
    if options.nevents < 1 or options.nruns < 1:
        sys.exit("Please specify a number > 1 for number of events/runs")

    if options.nruns == 1:
        main(args[0], options.nevents, options.cfg_file)
    else:
        import time
        import string
        for i in range(options.nruns):
            timestr = string.replace(time.asctime(time.localtime()), ' ', '-')
            filename = args[0] + '_' + timestr + '.root'
            print '-' * 65
            print 'Saving to file %s' % filename
            print '-' * 65

            main(filename, options.nevents, options.cfg_file)
